import numpy as np
from Utils import GradUtils
import Utils
from sympy import *

def dataGenerator():
    _x1, _x2 = symbols('x1 x2')
    _f_x = _x1**2 + _x2**2 + 6*_x1 + 9
    _x = np.array((2, 0)).reshape((2, 1))
    _A = np.array([[2,1],
                   [1,0],
                   [0,1]])
    _b = np.array([4, 0, 0]).reshape((3, 1))
    _E=np.empty((0,0))
    return {
        'f_x': _f_x,
        'A': _A,
        'E':_E,
        'x': _x,
        'b': _b,
    }

if __name__=='__main__':
    data=dataGenerator()
    f_x = data['f_x']
    A = data['A']
    b = data['b']
    x = data['x']
    l=symbols('l')
    cnt=0
    print('既约投影法演示')
    print('-' * 40)
    print('初始数据:')
    print('\t\t目标函数为：')
    print('\t\tf(x)={}'.format(f_x))
    print('\t\t非有效约束为的矩阵为\n{}'.format(np.append(A, b, axis=1)))
    print('-' * 40)
    while True:
        x=Utils.correctX(x)
        print('第{}次迭代，x={}，f(x)={}'.format(cnt, x.T, f_x.evalf(subs=Utils.list2dict(x))))
        print('-'*40)
        cnt=cnt+1
        division=GradUtils.saperateMatrix(A,x,b)
        A1=division['A1']
        A2=division['A2']
        b2=division['b2']
        P=GradUtils.genProjection(A1)
        delta_f=GradUtils.getGradient(f_x, x)
        d=-np.matmul(P,delta_f.reshape(len(delta_f),1))
        d=Utils.correctX(d)
        if np.all(d==0):
            CA1=GradUtils.correctA1(A1,delta_f)
            if CA1['KT']:
                break
            A1=CA1['A1']
            P=GradUtils.genProjection(A1)
            d = -np.matmul(P, delta_f.reshape(len(delta_f), 1))
        maxLambda=GradUtils.getMaxLambda(A2, b2, x, d)
        xl=x+l*d
        xl=Utils.list2dict(xl)
        f_l=f_x.evalf(subs=xl)
        x=x+Utils.getMinPoint(f_l,maxLambda,l)*d
    print('找到KT点：x={},此时函数最小值为{}'.format(x.reshape((1,-1)),f_x.evalf(subs=Utils.list2dict(x))))
